﻿#light
namespace AVL

module AVLTree =
    (*The AVLTree type declaration. Self-explanatory.
      The short stores the height of the tree;
      This tree will stop working properly
      should it contain more than 2^32767 elements
      which is a very unrealistic size.*)
    type 'a AVLTree when 'a :> System.IComparable (* <- make sure elements can be sorted *) =
        | Node of 'a * 'a AVLTree * 'a AVLTree * int16
        | Empty

    //Gets the value of the tree and its two subtrees as a tuple
    let root = function
        | Node(root,leftChild,rightChild,_) -> root,leftChild,rightChild
        | Empty -> failwith("Attempt was made to get data from an empty tree")
        
    //Gets the left subtree
    let left tr = let _,x,_ = root tr in x

    //Gets the right subtree
    let right tr = let _,_,x = root tr in x

    //These functions are safe versions of left and right
    //If you call them on an empty tree, instead of an exeception, an empty tree is returned
    let sleft = function
        | Empty -> Empty
        | x -> left x
    let sright = function
        | Empty -> Empty
        | x -> right x

    //Gets the value at the root of the tree
    let rootValue tr = let x,_,_ = root(tr) in x

    //Checks if a tree is the empty tree
    let isEmpty = function
        | Empty -> true
        | _ -> false
        
    //Gets the minimum of a tree
    let rec min = function
    | Empty -> failwith "Tried to find minimum of an empty tree"
    | Node(x,l,r,_) when isEmpty l -> x
    | t -> min (left t)
    
    //Gets the maximum of a tree
    let rec max = function
    | Empty -> failwith "Tried to find maximum of an empty tree"
    | Node(x,l,r,_) when isEmpty r -> x
    | t -> max (right t)

    //Retrieves the in-order traversal as a lazy sequence
    //Note that the deep recursion is bound by log(n) height!
    //Stack will not overflow in most contexts
    let rec toSeq tr =
        if isEmpty tr then Seq.empty
        else
            seq{
                yield! toSeq(left(tr));
                yield  rootValue(tr);
                yield! toSeq(right(tr));
               }
           
    //Retrieves the in-order traversal as an eager list
    //Will not compile on older versions, use the deprecated
    //Seq.to_list instead
    let toList tr = toSeq tr |> Seq.toList

    //This retrieves the stored height
    let height = function
            | Empty -> 0s
            | Node(_,_,_,h) -> h
            
    //The empty tree
    let empty = Empty

    //This function makes a tree with the proper stored height
    let makeTree root left right =
        let myheight = 1s + (if ((height left) > (height right)) then height left else height right)
        Node(root,left,right,myheight)
        
    //Positive values means skewed to the left, negative means skewed to the right. See wikipedia
    let balanceFactor tr =
        if(isEmpty tr) then 0s
        else (height(left(tr))) - (height(right(tr)))
        
    //Checks whether or not the tree contains an element
    let rec contains v tr =
        if(isEmpty tr) then false
        else if(v = (rootValue tr)) then true
        else if(v > (rootValue tr)) then contains v (right tr)
        else contains v (left tr)

    //Values of the tree are immutable, use closures to represent partially completed trees
    //These functions will return a function that replaces one subtree of the given tree

    //Replaces the right subtree
    let reconstructLeft  tr = let t,l,_ = root(tr) in (fun r -> makeTree t l r)

    //Replaces the left subtree
    let reconstructRight tr = let t,_,r = root(tr) in (fun l -> makeTree t l r)

    //Tree rotations
    //If you are struggling to understand these, try drawing them out
    //Look at each binding made, and the final tree made with those bindings
    //They are taken straight from the examples in wikipedia
    let rotateLL tr =
        let T3 = right tr
        let C  = rootValue tr
        let T2 = right (left tr)
        let B  = rootValue (left tr)
        let A  = rootValue (left (left tr))
        let T1 = right     (left (left tr))
        let T0 = left      (left (left tr))
        makeTree B (makeTree A T0 T1) (makeTree C T2 T3)

    let rotateLR tr =
        let T3 = right tr
        let C  = rootValue tr
        let T0 = left (left tr)
        let A  = rootValue (left tr)
        let B  = rootValue (right (left tr))
        let T2 = right     (right (left tr))
        let T1 = left      (right (left tr))
        makeTree B (makeTree A T0 T1) (makeTree C T2 T3)
                
    let rotateRR tr =
        let T0 = left tr
        let A  = rootValue tr
        let T1 = left (right tr)
        let B  = rootValue (right tr)
        let C  = rootValue (right (right tr))
        let T3 = right     (right (right tr))
        let T2 = left      (right (right tr))
        makeTree B (makeTree A T0 T1) (makeTree C T2 T3)
                            
    let rotateRL tr =
        let T0 = left tr
        let A  = rootValue tr
        let T3 = right (right tr)
        let C  = rootValue (right tr)
        let B  = rootValue (left (right tr))
        let T2 = right     (left (right tr))
        let T1 = left      (left (right tr))
        makeTree B (makeTree A T0 T1) (makeTree C T2 T3)
        
    (* Tree balancing function. Constant time. 
       Does not recursively attempt to balance subtrees
       Reason for this is because this function gets composed
       during tree rebuilding, so it will already have been
       called on all the needed subtrees
    *)
    let rec balance tr =
        if(balanceFactor tr > 1s) then
            if(balanceFactor (left tr) < 0s) then
                rotateLR tr
            else
                rotateLL tr
        else if(balanceFactor tr < -1s) then
            if(balanceFactor(right tr) > 0s) then
                rotateRL tr
            else
                rotateRR tr
        else tr

    (* Mutation operations follow
       Please note, trees are immutable and cannot be modified in place
       Function composition is used to build up a function that will
       rebuild the tree when given the last needed subtree.
       Tree rebuilding is logarithmic in time and space, thus stack 
       overflows should be virtually impossible under most conditions.
    *)

    //Insertion                
    let insert v tr =
        let rec loop t (reconstructor: 'a AVLTree -> 'a AVLTree) =
            if(isEmpty t) then reconstructor(makeTree v Empty Empty)
            else
                let r = rootValue t
                if(v = r) then tr
                else if(v > r) then loop (right t) (reconstructor << balance << reconstructLeft t)
                else loop (left t) (reconstructor << balance << reconstructRight t)
        loop tr (fun x -> x)
        
    //Deletion with in-order predecessor
    // removes the root and replaces it with a predecessor
    let deleteViaLeft tr =
        let rightSubtree = right tr
        let rec loop tr reconstructor =
            match (right tr) with
            | Empty -> (rootValue tr),(reconstructor (sleft tr)) //This is a predecessor
            | _ -> loop (right tr) (reconstructor << balance << reconstructLeft tr) //This is not a predecessor
        let nroot,nleft = loop (left tr) (fun x -> x)
        makeTree nroot nleft rightSubtree

    //Deletion with in-order successor
    // removes the root and replaces it with a successor
    let deleteViaRight tr =
        let leftSubtree = left tr
        let rec loop tr reconstructor =
            match (left tr) with
            | Empty -> (rootValue tr),(reconstructor (sright tr)) //This is a successor
            | _ -> loop (left tr) (reconstructor << balance << reconstructRight tr) //This is not a successor
        let nroot,nright = loop (right tr) (fun x -> x)
        makeTree nroot leftSubtree nright
        
    //This deletion chooses the appropriate in-order neighbor depending on the height of the tree
    let rec delete v tr =
        let rec loop t reconstructor =
            if(isEmpty t) then tr
            else if(v = rootValue t) then
                if(isEmpty(left t) && isEmpty(right t)) then reconstructor(Empty)
                else if(isEmpty(left t)) then reconstructor(right t)
                else if(isEmpty(right t)) then reconstructor(left t)
                else if(balanceFactor t > 0s) then reconstructor (deleteViaLeft t)
                else reconstructor (deleteViaRight t)
            else if(v > rootValue t) then
                loop (right t) (reconstructor << balance << reconstructLeft t)
            else loop (left t) (reconstructor << balance << reconstructRight t)
        loop tr (fun x -> x)
    
    //Here ends the full implementation of the AVL tree. Feel free to submit comments. 
    //Ask "barendventer" questions if you have any.
    //Play with the code by all means, if you installed F# you will have access to FSI
    //FSI will help you learn how individual parts of the application work
    
    
    //Here is some code specifically for our GUI. May be best to refactor this into a seperate namespace
    //Code is quite messy from this point on
    //This removes duplicates from a list - better way to do this might be by using Seq.pairwise and Seq.collect
    let rec removeDups last = function
        | [] -> Seq.empty
        | h::t -> if h = last then seq{yield! removeDups last t} else seq{yield h; yield! removeDups h t}
    let callRemoveDups = function
        | [] -> Seq.empty
        | h::t -> seq{yield h; yield! removeDups h t}
        
    let rec insertIter v tr =
        let rec loop t (reconstructor: 'a AVLTree -> 'a AVLTree) (rebalancer: (('a AVLTree -> 'a AVLTree) * ('a AVLTree -> 'a AVLTree)) list) =
            if(isEmpty t) then 
                List.scan (fun (acc,_) (f,r) -> (f acc),r) ((makeTree v Empty Empty),reconstructor) rebalancer
                |> List.map (fun (t,r) -> r t)
                |> callRemoveDups
            else
                let r = rootValue t
                if(v = r) then Seq.singleton tr
                else if(v > r) then 
                    loop 
                        (right t) 
                        (reconstructor << (reconstructLeft t)) 
                        (((balance << (reconstructLeft t)),reconstructor)::rebalancer)
                else loop 
                        (left t) 
                        (reconstructor << (reconstructRight t)) 
                        (((balance << (reconstructRight t)),reconstructor)::rebalancer)
        loop tr id []
        
    (* This is meant to provide the animations for the GUI. The code gets very messy here
       If you're interested in some of the more interesting features of F#, I use some of them here
       I don't always do the best thing though, and the deletion in particular was a bit of a hack
       I was very short on time at the time I wrote it.
       Perhaps you could re-inmplement this in a much better way? *)
    //Deletion with in-order predecessor
    let deleteViaLeftIter reconstructor tr =
        let rec loop tr reconstructor rebalancer =
            match (right tr) with
            | Empty -> (rootValue tr),(sleft tr),reconstructor,rebalancer
            | _ -> loop (right tr) (reconstructor << reconstructLeft tr) (((balance << reconstructLeft tr),reconstructor)::rebalancer)
        let nroot,l,rebuildl,steps = loop (left tr) reconstructor [(balance << reconstructRight tr),reconstructor]
        let steptrees = 
            List.scan (fun (acc,_) (f,l) -> (f acc) , l) (l,rebuildl) steps 
            |> List.tail 
            |> List.map (fun (t,l) -> l t)
        let rec buildSeq nroot lastTree = function
            | [] -> Seq.singleton(makeTree nroot (sleft lastTree) (sright lastTree))
            | h::t -> if h = lastTree then seq{yield! buildSeq nroot lastTree t;}
                      else
                          seq{
                              yield h;
                              yield! buildSeq nroot h t;
                             }
        buildSeq nroot (List.head steptrees) (List.tail steptrees)
           
    //Deletion with in-order successor
    let deleteViaRightIter reconstructor tr =
        let rec loop tr reconstructor rebalancer =
            match (left tr) with
            | Empty -> (rootValue tr),(sright tr),reconstructor,rebalancer
            | _ -> loop (left tr) (reconstructor << reconstructRight tr) (((balance << reconstructRight tr),reconstructor)::rebalancer)
        let nroot,r,rebuildr,steps = loop (right tr) reconstructor [(balance << reconstructLeft tr),reconstructor]
        let steptrees = 
            List.scan (fun (acc,_) (f,r) -> (f acc) , r) (r,rebuildr) steps 
            |> List.tail 
            |> List.map (fun (t,r) -> r t)
        let rec buildSeq nroot lastTree = function
            | [] -> Seq.singleton(makeTree nroot (sleft lastTree) (sright lastTree))
            | h::t -> if h = lastTree then seq{yield! buildSeq nroot lastTree t;}
                      else
                          seq{
                              yield h;
                              yield! buildSeq nroot h t;
                             }
        buildSeq nroot (List.head steptrees) (List.tail steptrees)          
      
    //Deletion with a step iterator to animate individual steps taken
    let rec deleteIter v tr =
        let rec loop t reconstructor rebalancer =
            if(isEmpty t) then Seq.empty,reconstructor,Empty,[]
            else if(v = rootValue t) then
                if(isEmpty(left t) && isEmpty(right t)) then Seq.empty,reconstructor,Empty,rebalancer
                else if(isEmpty (left t)) then Seq.empty,reconstructor,(right t),rebalancer
                else if(isEmpty (right t)) then Seq.empty,reconstructor,(left t),rebalancer
                else if(balanceFactor t > 0s) then (deleteViaLeftIter reconstructor t),(reconstructor),(deleteViaLeft t),(rebalancer)
                else (deleteViaRightIter reconstructor t),reconstructor,(deleteViaRight t),rebalancer
            else if(v > rootValue t) then
                loop (right t) (reconstructor << reconstructLeft t) (((balance << reconstructLeft t),reconstructor)::rebalancer)
            else loop (left t) (reconstructor << reconstructRight t) (((balance << reconstructRight t),reconstructor)::rebalancer)
        let deletion,reconstructor,rerootedTree,balanceUp = loop tr id []
        let rec buildSeq lastTree = function
            | [] -> Seq.empty
            | h::t ->
                if h = lastTree then buildSeq lastTree t
                else seq{yield h; yield! buildSeq h t}
        let treeList = List.scan (fun (acc,_) (f,r) -> (f acc),r) (rerootedTree,reconstructor) balanceUp |> List.map (fun (t,r) -> r t)
        let outsq = Seq.append deletion (buildSeq rerootedTree treeList)
        if Seq.isEmpty outsq && not(isEmpty tr) then Seq.singleton (delete v tr)
        else outsq